import os
import sys
import pandas as pd
import numpy as np
import matplotlib  
matplotlib.use('Agg') 
import matplotlib.pyplot as plt
from matplotlib.pyplot import savefig  
import math
from matplotlib.colors import LogNorm
from scipy import interpolate
import matplotlib.patches as mpatches

#get distance, a=lat1, b=lat2, c=lng1, d=lng2
def get_distance(a,b,c,d):
	'''calculate the distance'''
	PI=math.pi/180
	R=6371.004
	C=np.sin(a*PI)*np.sin(b*PI)+np.cos(a*PI)*np.cos(b*PI)*np.cos((c-d)*PI)
	Distance=R*np.arccos(np.round(C,12))
	return Distance

def plot_hist(x,fname,b=None):
	fig = plt.figure(figsize=(12,12),dpi=200)
	ax = fig.add_subplot(111)
	ax.hist(x,bins=b)
	fig.savefig(fname,dpi=200)		
	plt.close(fig)

# @_savedata('temp_out.pkl')
def calc_potential(infiles,outfiles=[],params=[]):
# mids = range(19)
# mnames = ['HYC','DeJi','XinBai','ZY','DY','JL','AS','HX_WD','SYC','SPF','SLME','HuaYC','XL_JY','JN_WD','TXJR','JN_JY','JF','21Cent','XJK_JY']#,'YH']

# msel = [0,8,9]

# mids = range(10)
# mnames=['Mall%d'%mi for mi in range(20)]

	file_info,file_home = infiles
	out_file= outfiles[0]
	attri 	= eval(params[0])
	mids	= eval(params[1])
	mnames	= eval(params[2])
	mn_dict = dict(zip(mids,mnames))
	print 'mn_dict', mn_dict
	print 'attri',attri
	
	
	smin,smax = -9.,6.
	rint = 200
	tag = 't'
	bSize = False
	
	def prepData(file_info,file_home,attri):
		v = pd.read_csv(file_info)
		h = pd.read_csv(file_home,names=['mac','ci','y','x','unidays','dm2'])
		if attri!='all':
			v = v[v.day_attri==attri]

		v = v[v.MallID.isin(mids)]
		# if False:
			# for mid, g in v.groupby('MallID'):
				# q = g[['mac','date']].groupby('mac',as_index=False).agg(np.size)
				# q = pd.merge(q,h,on='mac')
				# plot_hist(q.date,'hist_mallVisitTimes_%d.jpg'%mid,range(20))
				# plot_hist(q.date/q.unidays.astype(float),'hist_timesRatio_%d.jpg'%mid,np.arange(0,0.5,0.05))

		v = pd.merge(v,h[['mac','ci','x','y']],on='mac',left_index=True).reset_index()
		return v,h
		
	def plot_pot(g,mid):
		# f = interpolate.interp2d(g.xm, g.ym, g.pot, kind='cubic')
		# xa = np.arange(118.5,119,0.005)
		# ya = np.arange(31.8,32.3,0.005)
		# x,y = np.meshgrid(xa,ya)
		# z = f(xa,ya)		
		fig = plt.figure(figsize=(8,8),dpi=200)
		ax = fig.add_subplot(111)
		# z=(g.pot-g.pot.min())/(g.pot.max()-g.pot.min())
		z=np.abs((g.score-smin)/(smax-smin))
		ax.scatter(g.x,g.y,s=2,marker='s',c=z,cmap=plt.cm.rainbow)
		#ax.contourf(g.x,g.y,z,cmap=plt.get_cmap('rainbow'))
		ax.set_xlim((118.5,119))
		ax.set_ylim((31.8,32.3))
		fig.savefig('pot_contour_mall%d.jpg'%mid,dpi=200)
		plt.close(fig)
	
	v,h = prepData(file_info,file_home,attri)
	pot={}
	v.loc[:,'pot'] = 0
	
	v.loc[:,'xm'] = np.round(v.x.values*rint)/rint
	v.loc[:,'ym'] = np.round(v.y.values*rint)/rint
	
	md=pd.DataFrame()
	ciall = pd.DataFrame()	
	for mid in mids:
		ids = v.index[v.MallID==mid]
		g = v.loc[ids,:]
		v.loc[ids,'pot'] = g.cost_time_sum*np.exp(-g.dist*g.dist/10)
		#v.loc[ids,'pot'] = np.exp(-g.dist*g.dist/10) #tested to see no big diff
		
		pot[mid] = np.sum(v.loc[ids,'pot'])
		g = v.loc[ids,:]
		g = g[['xm','ym','pot']].groupby(['xm','ym'],as_index=False).agg(np.sum)
		g.loc[:,'pot'] = np.log10(g.pot)
		g.rename(columns={'pot':'score','xm':'x','ym':'y'},inplace=True)
		
		# print g.score.describe()
		
		g.to_csv('mall%d_score_grid%d.csv'%(mid,rint),index=False)
		g = g[g.score>smin]
		plot_pot(g,mid)
		# md[mid] = g
		g.loc[:,'mid']=mid
		md = pd.concat([g,md])
		ciall = pd.concat([g[['x','y']],ciall])
	
	ciall.drop_duplicates(inplace=True)
	
	hxy = h[['ci','x','y']].drop_duplicates()
	hxy.loc[:,'x'] = np.round(hxy.x.values*rint)/rint
	hxy.loc[:,'y'] = np.round(hxy.y.values*rint)/rint
	hxy.drop_duplicates(['x','y'],inplace=True)
	hxy['xys'] = hxy.x.astype(str)+'_'+hxy.y.astype(str)
	
	hh = hxy[hxy.x.isin(ciall.x) & hxy.y.isin(ciall.y)]
	hh.sort_values('xys',inplace=True)
	
	# print 'potential - {mallid:pot_value}',pot
	md.to_csv(outfiles[0],index=False)
	hh.to_csv(outfiles[1],index=False)
	return md, hh